from abc import ABC, abstractmethod
from typing import Optional, List, Union

from langchain_core.messages import SystemMessage
from pydantic import BaseModel, Field, ConfigDict

from app.llm import LLM
from app.logger import logger
from app.token_counter import TokenCounter


class BaseAgent(BaseModel, ABC):  # 继承自 ABC 来定义抽象类
    """
    基础智能体：提供基本的状态管理和执行循环。
    为状态转换、内存管理、以及一个基于‘ step ’的执行循环。子类必须实现‘ step ’方法。
    """

    # Core attributes
    name: str = Field(..., description="Unique name of the agent")  # ... 作为 Field() 的第一个参数表示该字段是必需的。
    description: Optional[str] = Field(None, description="Optional agent description")  # None 表示该字段可选，默认值为 None

    # Prompts
    system_prompt: Optional[List[Union[dict, SystemMessage]]] = None

    # Dependencies
    llm: LLM = Field(default_factory=LLM,  # default_factory表示该agent实例将包含一个语言模型组件，默认会自动创建LLM实例。
                     description="Language model instance")
    token_counter: Optional[TokenCounter] = None

    model_config = ConfigDict(
        arbitrary_types_allowed=True,
        extra="allow"
    )

    def __init__(self, **args): # 双引号是前向引用。意思此时类未完全定义，但相信我，未来会有一个叫 BaseAgent 的类型。是处理循环依赖和自引用的标准做法。
        """Initialize agent with default settings if not provided."""
        super().__init__(**args)
        if self.llm is None or not isinstance(self.llm, LLM):
            self.llm = LLM(config_name=self.name.lower())
        self.token_counter = TokenCounter(self.llm.client)

    @abstractmethod  # 使用 abstractmethod 装饰器来定义抽象方法。
    async def step(self) -> str:
        """Execute a single step in the agent's workflow.

        Must be implemented by subclasses to define specific behavior.
        """

    async def run(self) -> list:
        # 通过 async with 语句使用异步资源管理功能。
        logger.info(f"Agent is running.")
        res = await self.step()
        logger.info(f"Agent completed.")
        return res
